import matplotlib.pyplot as plt
from sklearn.cluster import DBSCAN
import pandas as pd
import numpy as np
from General.MapLoader import MapLoader
import networkx as nx
from sklearn.cluster import KMeans
from sklearn.cluster import SpectralClustering
from scipy.cluster.hierarchy import linkage, dendrogram, fcluster


# pd.set_option('display.max_columns', None)
# pd.set_option('display.max_rows',None)

class NodeClusterHandler:
    def __init__(self, map_loader):
        self.map_loader = map_loader
        # self.cluster_result_output_filepath = '../data/node_index_with_cluster.xlsx'

    def plot_clusters(self, node_indices, labels, map_code):
        """
        绘制聚类结果
        :param node_indices: 节点索引
        :param labels: 聚类标签
        :param map_code: 楼层代码
        """
        # 获取节点坐标
        positions = {index: (self.map_loader.node_features[index]['node_x'],
                             self.map_loader.node_features[index]['node_y']) for index in node_indices}

        # 创建网络图
        G = nx.Graph()

        # 添加节点和边
        for start_index in node_indices:
            G.add_node(start_index, pos=positions[start_index])
            for end_index in node_indices:
                if start_index != end_index and self.map_loader.adjacency_matrix[start_index, end_index] != np.inf:
                    G.add_edge(start_index, end_index)

        # 设置节点颜色
        colors = ['gray' if label == -1 else plt.cm.jet(label / len(set(labels))) for label in labels]

        # 绘制图
        plt.figure(figsize=(10, 8))
        nx.draw_networkx(G, pos=positions, node_color=colors, with_labels=False, node_size=50)

        plt.title(f'Clusters for Map Code: {map_code}')
        plt.xlabel('X Coordinate')
        plt.ylabel('Y Coordinate')
        plt.axis('equal')
        plt.show()

    def perform_clustering_by_kmeans(self, map_code, n_clusters=5, output_file='node_index_with_cluster.xlsx'):
        """
        对指定楼层进行 K-Means 聚类，并将结果保存到 DataFrame 中，最后使用这个方法聚类
        :param map_code: 要聚类的楼层代码
        :param n_clusters: 聚类数量
        :param output_file: 输出文件名
        """
        df = pd.read_excel(output_file)
        graphLoader = self.map_loader
        adjacency_matrix = graphLoader.get_adj_matrix()
        adjacency_matrix = np.where(np.isinf(adjacency_matrix), 0, adjacency_matrix)

        # 筛选当前楼层的数据
        layer_data = df[df['map_code'] == map_code]
        indices = layer_data['index'].values

        # 获取节点坐标
        coordinates = layer_data[['node_x', 'node_y']].values

        # 应用 K-Means 聚类
        kmeans = KMeans(n_clusters=n_clusters, random_state=0)
        labels = kmeans.fit_predict(coordinates)

        # 计算当前最大标签值
        max_label = df['cluster'].max() if 'cluster' in df.columns else 0

        # 将聚类结果添加到 DataFrame，并递增当前标签
        df.loc[df['map_code'] == map_code, 'cluster'] = labels + max_label + 1  # +1 确保从1开始

        # 输出结果到新的文件
        # df.to_excel(output_file, index=False)
        # print(f"Clustering results saved to {output_file}")

        # 可视化
        G = nx.Graph()
        # 添加节点
        for i, (x, y) in enumerate(coordinates):
            node_index = layer_data.iloc[i]['index']
            G.add_node(node_index, pos=(x, y))

        # 添加边
        for i in range(len(indices)):
            for j in range(len(indices)):
                if adjacency_matrix[indices[i], indices[j]] != 0:  # 如果有连接
                    G.add_edge(layer_data.iloc[i]['index'], layer_data.iloc[j]['index'])

        pos = nx.get_node_attributes(G, 'pos')

        # 使用颜色映射
        colors = [labels[i] for i in range(len(labels))]
        color_map = plt.get_cmap("tab10")  # 选择一个具有明显区分的颜色映射

        plt.figure(figsize=(10, 8))
        nx.draw_networkx(G, pos, node_color=[color_map(c % 10) for c in colors], with_labels=False, node_size=30)

        # 标注质心
        # centroids = kmeans.cluster_centers_
        # for idx, (x, y) in enumerate(centroids):
        #     plt.scatter(x, y, c='red', marker='X', s=200, label=f'Centroid {idx + 1}')

        plt.title(f'K-Means Clustering for Map Code {int(map_code)}')
        plt.xlabel('X Coordinate')
        plt.ylabel('Y Coordinate')
        plt.grid()
        plt.grid(False)
        plt.axis('equal')
        plt.legend()
        plt.show()



if __name__ == '__main__':
    # 创建 MapLoader 实例
    map_loader = MapLoader()
    map_loader.create_graph_from_files()

    # 创建 Cluster 实例并进行聚类
    clusterer = NodeClusterHandler(map_loader)
    clusterer.perform_clustering_by_kmeans(map_code=21, n_clusters=20, output_file='../data/node_index_with_cluster.xlsx')




